/*************************************************************************
 @File Name: main.cc
 @Author: dzhang
 @Created Time : Sat 16 Jan 2021 01:05:42 PM CST
 @Description:

 ************************************************************************/

//  Created by Linzaer on 2019/11/15.
//  Copyright © 2019 Linzaer. All rights reserved.
#include <iostream>
#include <opencv2/opencv.hpp>
#include "monolith/interface/inference/ultra_face_dect_model.h"
#include "monolith/interface/inference/face_rec_model.h"
#include "monolith/interface/inference/common.h"
#include "monolith/interface/inference/utils.h"
#include "monolith/interface/sql/sql_face.h"
#include "glog/logging.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_split.h"
#include "monolith/util/base64.h"

using namespace std;
USING_NAME_SPACE

const int DectModelWidth = 320;
const int DectModelHeight = 240;

const int RecModelWidth = 112;
const int RecModelHeight = 112;

const int FaceFeatureSize = 128;

const float FaceDistanceTheshold = 0.9;

const std::string RootPath = "../";

int get_frames(std::vector<cv::Mat>& frame_vec) {
    for (int frame_tag = 1; frame_tag < 286; frame_tag++) {
        std::string img_path = absl::StrCat("/home/dzhang/ArcFace/image-",
                frame_tag++, ".jpg");
        auto frame = cv::imread(img_path);
        frame_vec.push_back(frame);
    }
    return 0;
}

int register_person() {
	std::vector<std::string> img_list = {"../rec/ChandlerBing.webp", "../rec/JenniferAniston.jpeg", "../rec/dzhang.jpg", "../rec/zxli.jpg"};
	for (auto img_path : img_list) {
		LOG(INFO)
				<< "************************ Register Person ******************************";
		std::vector<std::string> str_list = absl::StrSplit(img_path, '/');
		std::string person_name = str_list[str_list.size() - 1];
		str_list = absl::StrSplit(person_name, ".");
		person_name = str_list[0];
		auto frame = cv::imread(img_path);

		int frame_width = frame.cols;
		int frame_height = frame.rows;

		std::vector<Box> face_boxes;
		cv::Mat image;
		cv::resize(frame, image, cv::Size(DectModelWidth, DectModelHeight));
		ml::inference::UltraFaceDectModel::get_instatnce()->run(image.data,
				face_boxes, frame_width, frame_height, 3);

		for (auto face : face_boxes) {
			cv::Point pt1(face.lx, face.ly);
			cv::Point pt2(face.rx, face.ry);
			int lx = face.lx, ly = face.ly;
			int width = face.rx - face.lx;
			int height = face.ry - face.ly;

			cv::Mat face_img;
			int crop_state = ml::inference::crop_face_img(frame, face_img, lx,
					ly, width, height);

			if (not crop_state) {
				cv::resize(face_img, face_img,
						cv::Size(RecModelWidth, RecModelHeight));
				cv::imshow("WinFace", face_img);
				cv::waitKey(1);
				std::vector<float> face_feature;
				ml::inference::FaceRecModel::get_instatnce()->run(face_img.data,
						face_feature, RecModelWidth, RecModelHeight, 3);
				FaceInfo face_info;
				ml::sql::SQLFace::get_instance()->query_face(face_feature,
						FaceDistanceTheshold, face_info);
				if (face_info.name == "") {
					face_info.name = absl::StrCat("", person_name);
					if (true) {
						ml::sql::SQLFace::get_instance()->insert_face(
								face_feature, face_info);
                        LOG(INFO) << "Register: " << face_info.name ;   
					}
				} else {
					LOG(INFO) << "Rec: " << face_info.name << ", "
							<< face_info.dis;
					cv::putText(frame, face_info.name, pt2,
							cv::FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255));
				}
			}
			cv::rectangle(frame, pt1, pt2, cv::Scalar(0, 255, 0), 2);
		}
		cv::imshow("win", frame);
		cv::waitKey(0);
	}

	return 0;

}

int init_module(){
    // init model
    string face_dect_mnn_path = RootPath + "model/faceDect_slim_320.mnn";
    ml::inference::ModelParameter face_dect_model_parameter = { .model_path_ =
            face_dect_mnn_path, .name_ = "UltraFaceDectModel", .input_height_ =
            DectModelHeight, .input_width_ = DectModelWidth,
            .input_channel_ = 3, .mean_ = { 127, 127, 127 }, .var_ = { 1.0
                    / 128, 1.0 / 128, 1.0 / 128 }, .num_thread_ = 4 };
    ml::inference::UltraFaceDectModel::init(face_dect_model_parameter, 0.8);

    string face_rec_mnn_path = RootPath + "model/faceRec_112x112.mnn";
    ml::inference::ModelParameter face_rec_model_parameter = { .model_path_ =
            face_rec_mnn_path, .name_ = "FaceRec", .input_height_ =
            RecModelHeight, .input_width_ = RecModelWidth, .input_channel_ = 3,
            .mean_ = { 0, 0, 0 }, .var_ = { 1.0, 1.0, 1.0 }, .num_thread_ = 4 };
    ml::inference::FaceRecModel::init(face_rec_model_parameter);

    // init sql
    ml::sql::SQLFace::get_instance()->open(RootPath + "db/face.db");

    return 0;
}

int main(int argc, char **argv) {
    init_module();
    // register_person();

    double fps, begin_time;
    char fps_string[1024];
    std::vector<cv::Mat> frame_vec;
    get_frames(frame_vec);
    // cv::VideoCapture capture(0);
    cv::VideoCapture capture("/home/dzhang/001.mp4");
    // cv::VideoCapture capture("/home/dzhang/friends.flv");

    int frame_tag = 0;
    int face_tag = 6;
    while (true) {
        // LOG(INFO)<< "****************************************************";
        ++frame_tag;
        begin_time = (double) cv::getTickCount();
        cv::Mat frame;
        capture >> frame;
        // frame = frame_vec[frame_tag++];

        int frame_width = frame.cols;
        int frame_height = frame.rows;

        std::vector<Box> face_boxes;
        cv::Mat image;
        cv::resize(frame, image, cv::Size(DectModelWidth, DectModelHeight));
        ml::inference::UltraFaceDectModel::get_instatnce()->run(image.data, face_boxes, frame_width, frame_height, 3);

        for (auto face : face_boxes) {
            cv::Point pt1(face.lx, face.ly);
            cv::Point pt2(face.rx, face.ry);
            int lx = face.lx, ly = face.ly;
            int width = face.rx - face.lx;
            int height = face.ry - face.ly;

            cv::Mat face_img;
            int crop_state = ml::inference::crop_face_img(frame, face_img, lx, ly, width, height);

            if(not crop_state) {
                cv::resize(face_img, face_img,
                        cv::Size(RecModelWidth, RecModelHeight));
                // cv::imshow("WinFace", face_img);
                // cv::waitKey(1);
                std::vector<float> face_feature;
                ml::inference::FaceRecModel::get_instatnce()->run(face_img.data, face_feature, RecModelWidth,
                        RecModelHeight, 3);
                FaceInfo face_info;
                ml::sql::SQLFace::get_instance()->query_face(face_feature, FaceDistanceTheshold, face_info);
                if(face_info.name == "") {
                    face_info.name = absl::StrCat("NO.", face_tag);
                    if(false) {
                        ml::sql::SQLFace::get_instance()->insert_face(face_feature, face_info);
                        face_tag++;
                    }
                } else {
                    LOG(INFO) << "Rec: " << face_info.name << ", " << face_info.dis;
                    cv::putText(frame, face_info.name, pt2, cv::FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255));
                }
            }

            cv::rectangle(frame, pt1, pt2, cv::Scalar(0, 255, 0), 2);

        }

        // display FPS
        begin_time = ((double) cv::getTickCount() - begin_time)
        / cv::getTickFrequency();
        fps = 1.0 / begin_time;
        sprintf(fps_string, "%.2f, %d", fps, frame_tag);
        std::string fpsString("FPS:");
        fpsString += fps_string;
        putText(frame, fpsString, cv::Point(5, 20), cv::FONT_HERSHEY_SIMPLEX,
                0.5, cv::Scalar(0, 0, 255));
        cv::imshow("Win", frame);
        cv::waitKey(1);
    }

    return 0;
}

// int main(){
//     return 0;
// }